package lib;

import java.util.*;
import java.util.concurrent.*;

public class Matrix2D {
    //生成随机矩阵
    public static int[][] randomMatrix(int row, int col) {
        if (row < 0 || col < 0)
            throw new IllegalArgumentException("Number of rows or number of columns less than zero.");
        if (row == 0 || col == 0) return null;
        int[][] m = new int[row][col];

        ForkJoinPool pool = new ForkJoinPool();
        RandomMatrix2DTask masterTask = new RandomMatrix2DTask(m, 0, row, 0, col);
        pool.invoke(masterTask);

        return m;
    }

    //并行矩阵加法
    public static long[][] parallelMatAdd(int[][] matA, int[][] matB) {
        if (matA == null || matB == null) throw new NullPointerException();
        if (matA.length != matB.length || matA[0].length != matB[0].length)
            throw new IllegalArgumentException("Sizes of two matrices not equal.");
        long[][] resultMat = new long[matA.length][matA[0].length];

        ForkJoinPool pool = new ForkJoinPool();
        MatrixAddTask masterTask = new MatrixAddTask(resultMat, matA, matB, 0, matA.length, 0, matA[0].length);
        pool.invoke(masterTask);

        return resultMat;
    }

    //并行矩阵乘法
    public static double[][] parallelMatMul(int[][] matA, int[][] matB) {
        if (matA == null || matB == null) throw new NullPointerException();
        if (matA[0].length != matB.length)
            throw new IllegalArgumentException("nCol of matrix A not equals to nRow of matrix B.");
        double[][] resultMat = new double[matA.length][matB[0].length];

        ForkJoinPool pool = new ForkJoinPool();
        MatrixMulTask masterTask = new MatrixMulTask(resultMat, matA, matB, 0, resultMat.length, 0, resultMat[0].length);
        pool.invoke(masterTask);

        return resultMat;
    }
}

//用于并行创建随机矩阵的任务类
class RandomMatrix2DTask extends RecursiveAction {
    private final int[][] m;
    private final int rowStart;
    private final int colStart; //此任务对象处理矩阵的[rowStart,rowEnd-1]行，下同
    private final int rowEnd;
    private final int colEnd;
    private static final int THRESHOLD = 20;

    public RandomMatrix2DTask(int[][] m, int rowStart, int rowEnd, int colStart, int colEnd) {
        this.m = m;
        this.rowStart = rowStart;
        this.rowEnd = rowEnd;
        this.colStart = colStart;
        this.colEnd = colEnd;
    }

    public void compute() {
        //任务规模小，串行解决
        int nRow = rowEnd - rowStart;
        int nCol = colEnd - colStart;
        if (nRow <= THRESHOLD || nCol <= THRESHOLD) {
            Random r = new Random();
            for (int i = rowStart; i <= rowEnd - 1; i++) {
                for (int j = colStart; j <= colEnd - 1; j++) {
                    m[i][j] = r.nextInt();
                }
            }
            return;
        }
        //拆分
        int rowMid = (rowStart + rowEnd) / 2;
        int colMid = (colStart + colEnd) / 2;
        RandomMatrix2DTask branchTaskA = new RandomMatrix2DTask(m, rowStart, rowMid, colStart, colMid);
        RandomMatrix2DTask branchTaskB = new RandomMatrix2DTask(m, rowStart, rowMid, colMid, colEnd);
        RandomMatrix2DTask branchTaskC = new RandomMatrix2DTask(m, rowMid, rowEnd, colStart, colMid);
        RandomMatrix2DTask branchTaskD = new RandomMatrix2DTask(m, rowMid, rowEnd, colMid, colEnd);
        invokeAll(branchTaskA, branchTaskB, branchTaskC, branchTaskD);
    }
}

//用于实现并行矩阵加法的任务类
class MatrixAddTask extends RecursiveAction {
    private final int[][] matA;
    private final int[][] matB;
    private final long[][] resultMat;
    private final int rowStart;
    private final int colStart; //此任务对象处理矩阵的[rowStart,rowEnd-1]行，下同
    private final int rowEnd;
    private final int colEnd;
    private static final int THRESHOLD = 20;

    public MatrixAddTask(long[][] resultMat, int[][] matA, int[][] matB, int rowStart, int rowEnd, int colStart, int colEnd) {
        this.resultMat = resultMat;
        this.matA = matA;
        this.matB = matB;
        this.rowStart = rowStart;
        this.rowEnd = rowEnd;
        this.colStart = colStart;
        this.colEnd = colEnd;
    }

    public void compute() {
        if (rowEnd - rowStart <= THRESHOLD || colEnd - colStart <= THRESHOLD) {
            for (int i = rowStart; i <= rowEnd - 1; i++) {
                for (int j = colStart; j <= colEnd - 1; j++) {
                    resultMat[i][j] = matA[i][j] + matB[i][j];
                }
            }
            return;
        }

        int rowMid = (rowStart + rowEnd) / 2;
        int colMid = (colStart + colEnd) / 2;
        MatrixAddTask branchA = new MatrixAddTask(resultMat, matA, matB, rowStart, rowMid, colStart, colMid);
        MatrixAddTask branchB = new MatrixAddTask(resultMat, matA, matB, rowStart, rowMid, colMid, colEnd);
        MatrixAddTask branchC = new MatrixAddTask(resultMat, matA, matB, rowMid, rowEnd, colStart, colMid);
        MatrixAddTask branchD = new MatrixAddTask(resultMat, matA, matB, rowMid, rowEnd, colMid, colEnd);
        invokeAll(branchA, branchB, branchC, branchD);
    }
}

//用于实现并行矩阵乘法的任务类
class MatrixMulTask extends RecursiveAction {
    private final int[][] matA;
    private final int[][] matB;
    private final double[][] resultMat;
    private final int rowStart;
    private final int colStart; //此任务对象处理矩阵的[rowStart,rowEnd-1]行，下同
    private final int rowEnd;
    private final int colEnd;
    private static final int THRESHOLD = 5;

    public MatrixMulTask(double[][] resultMat, int[][] matA, int[][] matB, int rowStart, int rowEnd, int colStart, int colEnd) {
        this.resultMat = resultMat;
        this.matA = matA;
        this.matB = matB;
        this.rowStart = rowStart;
        this.rowEnd = rowEnd;
        this.colStart = colStart;
        this.colEnd = colEnd;
    }

    public void compute() {
        if (rowEnd - rowStart <= THRESHOLD || colEnd - colStart <= THRESHOLD) {
            for (int i = rowStart; i <= rowEnd - 1; i++) {
                for (int j = colStart; j <= colEnd - 1; j++) {
                    //计算resultMat[i][j]
                    double sum = 0;
                    for (int k = 0; k <= matA[i].length - 1; k++) {
                        sum += matA[i][k] * matB[k][j];
                    }
                    resultMat[i][j] = sum;
                }
            }
            return;
        }

        int rowMid = (rowStart + rowEnd) / 2;
        int colMid = (colStart + colEnd) / 2;
        MatrixMulTask branchA = new MatrixMulTask(resultMat, matA, matB, rowStart, rowMid, colStart, colMid);
        MatrixMulTask branchB = new MatrixMulTask(resultMat, matA, matB, rowStart, rowMid, colMid, colEnd);
        MatrixMulTask branchC = new MatrixMulTask(resultMat, matA, matB, rowMid, rowEnd, colStart, colMid);
        MatrixMulTask branchD = new MatrixMulTask(resultMat, matA, matB, rowMid, rowEnd, colMid, colEnd);
        invokeAll(branchA, branchB, branchC, branchD);
    }
}